/* tftp protocol implementation */

/* Copyright (C) 2012 David Zanetti
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; version 2 of the License.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License along
 *  with this program; if not, write to the Free Software Foundation, Inc.,
 *  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

#include <avr/io.h>
#include <avr/pgmspace.h>
#include <util/atomic.h>
#include <stdlib.h>
#include <string.h>
#include "global.h"
#include "tftp.h"
#include "encx24j600.h"
#include "udp.h"
#include "serial.h"
#include "ip.h"
#include "fw_update.h"

/* TFTP request opcodes. read/write are RRQ and WRQ in RFC */
#define TFTP_READ 1
#define TFTP_WRITE 2
#define TFTP_DATA 3
#define TFTP_ACK 4
#define TFTP_ERROR 5

/* TFTP error codes from RFC1350 */
#define TFTP_ERR_UNDEF 0
#define TFTP_ERR_NOTFOUND 1
#define TFTP_ERR_ACCESS 2
#define TFTP_ERR_FULL 3
#define TFTP_ERR_ILLEGAL 4
#define TFTP_ERR_UNKNOWN 5
#define TFTP_ERR_EXISTS 6
#define TFTP_ERR_NOUSER 7

/* tftp is so simple we don't bother with a struct for it's header */

/* since we only ever handle one connection at a time, this is all the state we want */

/* the last block we've successfully acknowledged */
uint16_t last_block = 0;
/* request address, so someone can't send us requests out of the blue */
uint8_t request_addr[4] = {0,0,0,0};
/* counter for timeout of RX, 0xff is inactive */
uint8_t timeout;

/* init */
void tftp_init(void) {
	/* reset our various states */
	memset(request_addr,0,4);
	last_block = 0;
	timeout = 0xff;
}

/* Process incoming TFTP packets. */

uint8_t tftp_process(uint8_t slot,uint8_t *src,uint16_t len) {
	uint16_t opcode;
	uint16_t origlen;

	/* FIXME: this test is a bit crap. we *should* generate a unique source port for our outbound tftp request, but we don't */
	if (timeout != 0xff && memcmp(request_addr,src,4) != 0) {
		serial0_tx_PGM(PSTR("tftp: unexpected packet"));
		return E_NOTHING;
	}

	if (encx24j600_rx(&opcode,sizeof(opcode)) != sizeof(opcode)) {
		serial0_tx_PGM(PSTR("tftp: packet underrun\r\n"));
		return E_NOTHING;
	}

	/* put it in machine order */
	opcode = swap16(opcode);

	origlen = len;
	len -= 2; /* opcode field */

	switch (opcode) {
		case TFTP_ERROR:
		{
			uint16_t error_code;

			if (encx24j600_rx(&error_code,sizeof(error_code)) != sizeof(error_code)) {
				serial0_tx_PGM(PSTR("tftp: packet underrun\r\n"));
				return E_NOTHING;
			}
			/* could be for a variety of reasons, so just kill the connection */
			tftp_init();
			serial0_tx_PGM(PSTR("tftp: err code "));
			serial0_tx_dec(swap16(error_code));
			serial0_tx_cr();
			return E_NOTHING;
		}
			break; /* not reached */
		case TFTP_DATA:
			/* a data packet, see if this is the one we were expecting */
			{
				uint16_t blockid;

				if (encx24j600_rx(&blockid,sizeof(blockid)) != sizeof(blockid)) {
					serial0_tx_PGM(PSTR("tftp: packet underrun\r\n"));
					return E_NOTHING;
				}

				blockid = swap16(blockid); /* put in machine order */

				if (blockid == last_block+1) {

					/* compute packet length from UDP header */
					len -= 2; /* 2 bytes in headers for a data packet */

					/* copy the rest of the packet data, and then ack it */
					/* FIXME: this could be used for other things, but for now, we only fetch firmware updates */
					/* if we are in the first block, then we expect to see the following sequence: NTP14<crcmsb><crclsb> as the first bytes */
					if (blockid == 1) {
						uint8_t sig[7];
						/* read the signature and CRC */
						if (encx24j600_rx(&sig,sizeof(sig)) != sizeof(sig)) {
							serial0_tx_PGM(PSTR("tftp: packet underrun\t\n"));
							return E_NOTHING;
						}
						if (strncmp((char *)sig,"NTP14",5) != 0) {
							uint8_t tx[4];
							serial0_tx_PGM(PSTR("tftp: not a valid firmware update file\r\n"));
							/* report an error gracefully anyway */
							tx[0] = (TFTP_ERROR >> 8);
							tx[1] = (TFTP_ERROR & 0xff);
							tx[2] = (0);
							tx[3] = (0);
							if (!encx24j600_tx_prepend(slot,tx,4)) {
								serial0_tx_PGM(PSTR("tftp: failed to send nack\r\n"));
								tftp_init();
								return E_NOTHING;
							}
							return E_REPLY;
						}
						/* present CRC to the flash update interface */
						serial0_tx_PGM(PSTR("tftp: crc is "));
						serial0_tx_hex(sig[5]);
						serial0_tx_hex(sig[6]);
						serial0_tx_cr();
						fw_new(((uint16_t)sig[5] << 8) | (uint16_t)sig[6]);
						/* the rest of the packet may be processed as normal */
						/* adjust the length for further copies */
						len -= 7;
					}

					/* copy the packet data into the firmware update buffer */
					{
						/* FIXME: really, we should copy this data inside fw_update somewhere, instead of twice here */
						//uint8_t buf[512]; /* worst case for the payload */

						if (len > 512) {
							/* this is an invalid packet */
							serial0_tx_PGM(PSTR("tftp: invalid packet\r\n"));
							fw_discard();
						}

						/* pass to fw update process */
						fw_temp_write_ethernet(len); /* copy directly from packet */
					}

					/* generate ACK for this block */
					uint8_t tx[4];
					tx[0] = (TFTP_ACK >> 8);
					tx[1] = (TFTP_ACK & 0xff);
					tx[2] = (blockid >> 8);
					tx[3] = (blockid & 0xff);
					if (!encx24j600_tx_prepend(slot,tx,4)) {
						serial0_tx_PGM(PSTR("tftp: failed to send ack\r\n"));
						tftp_init();
						return E_NOTHING;
					}
					last_block = blockid;
					timeout = 20; /* 20 second timeout */

					/* test if this is the last packet, to kill the connection */
					if (origlen != 516) {
						serial0_tx_PGM(PSTR("tftp: transfer complete\r\n"));
						tftp_init();
						/* mark main loop we have firmware to install */
						set_flag(flag_stat,FLAG_STAT_FW_UPGRADE);
					}

					return E_REPLY;
				} else {
					serial0_tx_PGM(PSTR("tftp: out of order data packet\r\n"));
					timeout = 20; /* give us another 20 seconds */
					return E_NOTHING;
				}
			}
	}

	return E_NOTHING;
}

/* request a file from the given address and filename */
void tftp_request(uint8_t *addr,const char *filename) {
	char tx[256];
	uint8_t fn_len;
	/* initalise our state */
	last_block = 0;
	/* copy the source address we were given into the stack */
	memcpy(request_addr,addr,4);

	/* clear the TX packet */
	memset(tx,0,256);

	/* prepare our output packet */

	/* opcode */
	tx[0] = (TFTP_READ >> 8);
	tx[1] = (TFTP_READ & 0xff);
	/* stringcopy in the filename, up to 64 bytes */
	fn_len = strlen(filename);
	if (fn_len > 63) {
		fn_len = 63;
	};
	strncpy(tx+2,filename,fn_len);
	/* copy in the text "octet" for mode */
	strncpy(tx+2+fn_len+1,"octet",5);

	/* send the packet */
	if (!udp_send(addr,69,69,(uint8_t *)tx,2+fn_len+7)) {
		serial0_tx_PGM(PSTR("tftp: send failed\r\n"));
		timeout = 0xff;
		memset(request_addr,0,4);
		return;
	} else {
		serial0_tx_PGM(PSTR("tftp: requesting "));
		serial0_tx(filename,strlen(filename));
		serial0_tx_PGM(PSTR(" from "));
		ip_printaddr(addr);
		serial0_tx_cr();
	}

	return;
}

void tftp_tick(void) {
	/* if we have a counter to decrement, then do so */
	if (timeout != 0xff) {
		timeout--;
		if (timeout == 0) {
			serial0_tx_PGM(PSTR("tftp: abandoning connection\r\n"));
			/* give up and clear state */
			last_block = 0;
			memset(request_addr,0,4);
			timeout = 0xff;
		}
	}
	return;
}

